'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: agent_queue_manager.py
* @Time: 2025/11/17
* @All Rights Reserve By Brtc
'''
import queue
import time
import uuid
from queue import Queue
from typing import Generator
from uuid import UUID

from redis import Redis

from internal.entity.conversation_entity import InvokeFrom
from internal.entity.queue_entity import AgentQueueEvent, QueueEvent


class PublishFrom:
    """应用的发布来源"""
    APPLICATION_MANAGER = 1
    TASK_PIPELINE  =2


class AgentQueueManager:
    """只能队列管理"""
    q:Queue
    user_id:UUID
    task_id:UUID
    invoke_from:InvokeFrom
    redis_client:Redis

    def __init__(self, user_id:UUID,
                 task_id:UUID,
                 invoke_from:InvokeFrom,
                 redis_client:Redis)->None:
        """构造函数，初始化队列管理器"""
        #1、初始化数据
        self.q = Queue()
        self.task_id = task_id
        self.user_id = user_id
        self.invoke_from = invoke_from
        self.redis_client = redis_client
        #2、判断用户的是的类型， 目前有账号和网页(Debugger&WebApp)调用以及终端用户end_point
        user_prefix = "account" if invoke_from in [InvokeFrom.WEB_APP, InvokeFrom.DEBUGGER] else "end-user"
        #3、设置任务对应的缓存键， 代表此次任务开始
        self.redis_client.setex(
            self._generate_task_belong_cache_key(task_id),
            1800,
            f"{user_prefix}-{str(user_id)}"
        )

    def listen(self)->Generator:
        """监听队列函数返回生成的数据"""
        #1、定义基础数据记录超时间， 开始时间、最后一次ping 同时间
        listen_timeout = 600
        start_time = time.time()
        last_ping_time = 0
        #2、创建循环列表执行死循环监听队列数据
        while True:
            try:
                #3、从队列中提取数据检测数据是否存在, 如果存在则直接返回
                item = self.q.get(timeout=1)
                if item is None:
                    break
                yield item
            except queue.Empty:
                continue
            finally:
                #4、计算获取数据的总耗时，并且判断是否超时， 如果超时则往 队列添加停止事件
                elapsed_time = time.time() - start_time
                if elapsed_time > listen_timeout:
                    pass
                #5、每10s 发送一个 ping 请求检查客户端是否主动断开连接
                if elapsed_time//10 > last_ping_time:
                    last_ping_time = elapsed_time//10
                    self.publish(AgentQueueEvent(
                        id = uuid.uuid4(),
                        task_id=self.task_id,
                        event = QueueEvent.PING
                    ))
                #6、判断总耗时是否超时， 如果超时则添加超时事件
                if elapsed_time  >= listen_timeout:
                    self.publish(AgentQueueEvent(
                        id = uuid.uuid4(),
                        task_id=self.task_id,
                        event=QueueEvent.TIMEOUT
                    ))
                #7、检测是否停止， 如果已经停止则添加停止事件
                if self._is_stopped():
                    self.publish(AgentQueueEvent(
                        id = uuid.uuid4(),
                        task_id=self.task_id,
                        event = QueueEvent.STOP
                    ))




    def stop_listen(self)->None:
        """停止监听队列"""
        self.q.put(None)


    def publish(self, agent_queue_event:AgentQueueEvent)->None:
        """发布事件到队列中"""
        self.q.put(agent_queue_event)


    def publish_error(self, error)->None:
        """发布错误时间到队列中"""
        self.publish(AgentQueueEvent(
            id = uuid.uuid4(),
            task_id=self.task_id,
            event=QueueEvent.ERROR,
            observation=str(error)
        ))


    def _is_stopped(self)->bool:
        """检查任务是否停止"""
        task_stopped_cache_key = self._generate_task_stopped_cache_key(self.task_id)
        result = self.redis_client.get(task_stopped_cache_key)
        if result is not None:
            return True
        return False


    def set_stop_flask(self, task_id:UUID, invoke_from:InvokeFrom, user_id:UUID)->None:
        """设置任务停止标志"""
        #1、检查当前任务是否存在， 如果不存在则直接结束
        result:bytes = self.redis_client.get(self._generate_task_belong_cache_key(self.task_id))
        if result is None:
            return
        #2、查询当前任务和用户的信息是否匹配， 如果不匹配则直接结束
        user_prefix = "account" if invoke_from in [InvokeFrom.WEB_APP, InvokeFrom.DEBUGGER] else "end-user"
        #3、构建停止缓存键， 并设置缓存记录， 时间为600s
        task_stopped_cache_key = self._generate_task_stopped_cache_key(self.task_id)
        self.redis_client.setex(task_stopped_cache_key, 600, 1)


    @classmethod
    def _generate_task_belong_cache_key(cls, task_id:UUID)->str:
        """生成任务专属的缓存键"""
        return f"generate_task_belong:{str(task_id)}"


    @classmethod
    def _generate_task_stopped_cache_key(cls, task_id:UUID)->str:
        """生成任务停止缓存键"""
        return f"generate_task_stopped:{str(task_id)}"


